#include <stdio.h>
#include <math.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include "debug_utils.h"

// Calculate norm using cublasSdot
float calculateNormCublas(float* d_array, int size) {
    cudaError_t err;
    cublasStatus_t status;
    
    // Initialize cuBLAS
    cublasHandle_t handle;
    status = cublasCreate(&handle);
    if (status != CUBLAS_STATUS_SUCCESS) {
        printf("Failed to initialize cuBLAS: %d\n", status);
        return -1.0f;
    }
    
    // Calculate dot product of the vector with itself using cublasSdot
    float dot_result = 0.0f;
    status = cublasSdot(handle, size, d_array, 1, d_array, 1, &dot_result);
    if (status != CUBLAS_STATUS_SUCCESS) {
        printf("cublasSdot failed: %d\n", status);
        cublasDestroy(handle);
        return -1.0f;
    }
    
    // Wait for computation to finish
    err = cudaDeviceSynchronize();
    if (err != cudaSuccess) {
        printf("CUDA sync error: %s\n", cudaGetErrorString(err));
        cublasDestroy(handle);
        return -1.0f;
    }
    
    // Clean up and return result
    cublasDestroy(handle);
    return sqrtf(dot_result);
}

// Simple CUDA kernel to calculate sum of squares for norm (fallback implementation)
__global__ void sumOfSquaresKernel(float* input, float* output, int size) {
    extern __shared__ float sdata[];
    
    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Load into shared memory with bounds check
    float val = 0.0f;
    if (i < size) {
        val = input[i] * input[i]; // Square the value
    }
    sdata[tid] = val;
    __syncthreads();
    
    // Reduction in shared memory
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // Write result for this block to output
    if (tid == 0) {
        output[blockIdx.x] = sdata[0];
    }
}

// Helper function to calculate norm using custom CUDA kernel (fallback implementation)
float calculateNormCuda(float* d_array, int size) {
    cudaError_t err;
    
    // Determine grid and block sizes
    int blockSize = 256;
    int numBlocks = (size + blockSize - 1) / blockSize;
    
    // Allocate device memory for partial sums
    float* d_partial_sums = NULL;
    err = cudaMalloc((void**)&d_partial_sums, numBlocks * sizeof(float));
    if (err != cudaSuccess) {
        printf("CUDA malloc failed for partial sums: %s\n", cudaGetErrorString(err));
        return -1.0f;
    }
    
    // Launch kernel to compute partial sums
    sumOfSquaresKernel<<<numBlocks, blockSize, blockSize * sizeof(float)>>>(d_array, d_partial_sums, size);
    
    // Check for kernel errors
    err = cudaGetLastError();
    if (err != cudaSuccess) {
        printf("CUDA kernel launch error: %s\n", cudaGetErrorString(err));
        cudaFree(d_partial_sums);
        return -1.0f;
    }
    
    // Wait for kernel to finish
    err = cudaDeviceSynchronize();
    if (err != cudaSuccess) {
        printf("CUDA sync error: %s\n", cudaGetErrorString(err));
        cudaFree(d_partial_sums);
        return -1.0f;
    }
    
    // Copy partial sums to host
    float* h_partial_sums = (float*)malloc(numBlocks * sizeof(float));
    if (!h_partial_sums) {
        printf("Host memory allocation failed\n");
        cudaFree(d_partial_sums);
        return -1.0f;
    }
    
    err = cudaMemcpy(h_partial_sums, d_partial_sums, numBlocks * sizeof(float), cudaMemcpyDeviceToHost);
    if (err != cudaSuccess) {
        printf("CUDA memcpy failed: %s\n", cudaGetErrorString(err));
        free(h_partial_sums);
        cudaFree(d_partial_sums);
        return -1.0f;
    }
    
    // Sum the partial sums on CPU
    float sum_of_squares = 0.0f;
    for (int i = 0; i < numBlocks; i++) {
        sum_of_squares += h_partial_sums[i];
    }
    
    // Clean up
    free(h_partial_sums);
    cudaFree(d_partial_sums);
    
    // Return the square root of the sum
    return sqrtf(sum_of_squares);
}

// CPU fallback implementation when CUDA fails
float calculateNormCPU(float* d_array, int size) {
    float* h_array = (float*)malloc(size * sizeof(float));
    if (!h_array) {
        fprintf(stderr, "Host memory allocation failed\n");
        return -1.0f;
    }
    
    cudaError_t err = cudaMemcpy(h_array, d_array, size * sizeof(float), cudaMemcpyDeviceToHost);
    if (err != cudaSuccess) {
        fprintf(stderr, "Error copying data to host: %s\n", cudaGetErrorString(err));
        free(h_array);
        return -1.0f;
    }
    
    // Calculate norm on CPU
    float sum_of_squares = 0.0f;
    for (int i = 0; i < size; i++) {
        sum_of_squares += h_array[i] * h_array[i];
    }
    
    free(h_array);
    return sqrtf(sum_of_squares);
}

void debug_print_norm_cuda(const char* classname, const char* TAG, float* array, int size, int rank, int step) {
    if (rank == 0) {
        // First check if array is valid
        if (array == NULL) {
            printf("[%s] %s step=%d ERROR: Null pointer passed to debug_print_norm_cuda\n", 
                   classname, TAG, step);
            return;
        }
        
        // Make sure pending operations are complete
        cudaError_t err = cudaDeviceSynchronize();
        if (err != cudaSuccess) {
            printf("[%s] %s step=%d ERROR: CUDA sync failed: %s\n", 
                   classname, TAG, step, cudaGetErrorString(err));
            return;
        }
        
        // Check if the pointer is valid device memory
        cudaPointerAttributes attrs;
        err = cudaPointerGetAttributes(&attrs, array);
        if (err != cudaSuccess) {
            printf("[%s] %s step=%d ERROR: Invalid memory pointer: %s\n", 
                   classname, TAG, step, cudaGetErrorString(err));
            cudaGetLastError(); // Clear the error state
            return;
        }
        
        // Try to calculate the norm using cublasSdot first
        float norm = calculateNormCublas(array, size);
        
        // If cublasSdot fails, try the custom CUDA kernel implementation
        if (norm < 0) {
            printf("[%s] %s step=%d WARNING: cublasSdot norm calculation failed, trying custom CUDA implementation\n", 
                  classname, TAG, step);
            norm = calculateNormCuda(array, size);
            
            // If the custom CUDA implementation also fails, use the CPU fallback
            if (norm < 0) {
                printf("[%s] %s step=%d WARNING: Custom CUDA norm calculation failed, using CPU fallback\n", 
                      classname, TAG, step);
                norm = calculateNormCPU(array, size);
                
                if (norm < 0) {
                    printf("[%s] %s step=%d ERROR: All norm calculation methods failed\n", 
                          classname, TAG, step);
                    return;
                }
            }
        }
        
        // Print the result
        printf("[%s] %s step=%d norm=%.5E\n", classname, TAG, step, norm);
        fflush(stdout);
    }
}

// Calculate norm using cublasDdot for double precision
double calculateNormCublasDouble(double* d_array, int size) {
    cudaError_t err;
    cublasStatus_t status;
    
    // Initialize cuBLAS
    cublasHandle_t handle;
    status = cublasCreate(&handle);
    if (status != CUBLAS_STATUS_SUCCESS) {
        printf("Failed to initialize cuBLAS: %d\n", status);
        return -1.0;
    }
    
    // Calculate dot product of the vector with itself using cublasDdot
    double dot_result = 0.0;
    status = cublasDdot(handle, size, d_array, 1, d_array, 1, &dot_result);
    if (status != CUBLAS_STATUS_SUCCESS) {
        printf("cublasDdot failed: %d\n", status);
        cublasDestroy(handle);
        return -1.0;
    }
    
    // Wait for computation to finish
    err = cudaDeviceSynchronize();
    if (err != cudaSuccess) {
        printf("CUDA sync error: %s\n", cudaGetErrorString(err));
        cublasDestroy(handle);
        return -1.0;
    }
    
    // Clean up and return result
    cublasDestroy(handle);
    return sqrt(dot_result);
}

// Simple CUDA kernel to calculate sum of squares for norm (fallback implementation) for double precision
__global__ void sumOfSquaresKernelDouble(double* input, double* output, int size) {
    extern __shared__ double sdata_d[];
    
    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Load into shared memory with bounds check
    double val = 0.0;
    if (i < size) {
        val = input[i] * input[i]; // Square the value
    }
    sdata_d[tid] = val;
    __syncthreads();
    
    // Reduction in shared memory
    for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata_d[tid] += sdata_d[tid + s];
        }
        __syncthreads();
    }
    
    // Write result for this block to output
    if (tid == 0) {
        output[blockIdx.x] = sdata_d[0];
    }
}

// Helper function to calculate norm using custom CUDA kernel (fallback implementation) for double precision
double calculateNormCudaDouble(double* d_array, int size) {
    cudaError_t err;
    
    // Determine grid and block sizes
    int blockSize = 256;
    int numBlocks = (size + blockSize - 1) / blockSize;
    
    // Allocate device memory for partial sums
    double* d_partial_sums = NULL;
    err = cudaMalloc((void**)&d_partial_sums, numBlocks * sizeof(double));
    if (err != cudaSuccess) {
        printf("CUDA malloc failed for partial sums: %s\n", cudaGetErrorString(err));
        return -1.0;
    }
    
    // Launch kernel to compute partial sums
    sumOfSquaresKernelDouble<<<numBlocks, blockSize, blockSize * sizeof(double)>>>(d_array, d_partial_sums, size);
    
    // Check for kernel errors
    err = cudaGetLastError();
    if (err != cudaSuccess) {
        printf("CUDA kernel launch error: %s\n", cudaGetErrorString(err));
        cudaFree(d_partial_sums);
        return -1.0;
    }
    
    // Wait for kernel to finish
    err = cudaDeviceSynchronize();
    if (err != cudaSuccess) {
        printf("CUDA sync error: %s\n", cudaGetErrorString(err));
        cudaFree(d_partial_sums);
        return -1.0;
    }
    
    // Copy partial sums to host
    double* h_partial_sums = (double*)malloc(numBlocks * sizeof(double));
    if (!h_partial_sums) {
        printf("Host memory allocation failed\n");
        cudaFree(d_partial_sums);
        return -1.0;
    }
    
    err = cudaMemcpy(h_partial_sums, d_partial_sums, numBlocks * sizeof(double), cudaMemcpyDeviceToHost);
    if (err != cudaSuccess) {
        printf("CUDA memcpy failed: %s\n", cudaGetErrorString(err));
        free(h_partial_sums);
        cudaFree(d_partial_sums);
        return -1.0;
    }
    
    // Sum the partial sums on CPU
    double sum_of_squares = 0.0;
    for (int i = 0; i < numBlocks; i++) {
        sum_of_squares += h_partial_sums[i];
    }
    
    // Clean up
    free(h_partial_sums);
    cudaFree(d_partial_sums);
    
    // Return the square root of the sum
    return sqrt(sum_of_squares);
}

// CPU fallback implementation when CUDA fails for double precision
double calculateNormCPUDouble(double* d_array, int size) {
    double* h_array = (double*)malloc(size * sizeof(double));
    if (!h_array) {
        fprintf(stderr, "Host memory allocation failed\n");
        return -1.0;
    }
    
    cudaError_t err = cudaMemcpy(h_array, d_array, size * sizeof(double), cudaMemcpyDeviceToHost);
    if (err != cudaSuccess) {
        fprintf(stderr, "Error copying data to host: %s\n", cudaGetErrorString(err));
        free(h_array);
        return -1.0;
    }
    
    // Calculate norm on CPU
    double sum_of_squares = 0.0;
    for (int i = 0; i < size; i++) {
        sum_of_squares += h_array[i] * h_array[i];
    }
    
    free(h_array);
    return sqrt(sum_of_squares);
}

void debug_print_norm_cuda_d(const char* classname, const char* TAG, double* array, int size, int rank, int step) {
    if (rank == 0) {
        // First check if array is valid
        if (array == NULL) {
            printf("[%s] %s step=%d ERROR: Null pointer passed to debug_print_norm_cuda_d\n", 
                   classname, TAG, step);
            return;
        }
        
        // Make sure pending operations are complete
        cudaError_t err = cudaDeviceSynchronize();
        if (err != cudaSuccess) {
            printf("[%s] %s step=%d ERROR: CUDA sync failed: %s\n", 
                   classname, TAG, step, cudaGetErrorString(err));
            return;
        }
        
        // Check if the pointer is valid device memory
        cudaPointerAttributes attrs;
        err = cudaPointerGetAttributes(&attrs, array);
        if (err != cudaSuccess) {
            printf("[%s] %s step=%d ERROR: Invalid memory pointer: %s\n", 
                   classname, TAG, step, cudaGetErrorString(err));
            cudaGetLastError(); // Clear the error state
            return;
        }
        
        // Try to calculate the norm using cublasDdot first
        double norm = calculateNormCublasDouble(array, size);
        
        // If cublasDdot fails, try the custom CUDA kernel implementation
        if (norm < 0) {
            printf("[%s] %s step=%d WARNING: cublasDdot norm calculation failed, trying custom CUDA implementation\n", 
                  classname, TAG, step);
            norm = calculateNormCudaDouble(array, size);
            
            // If the custom CUDA implementation also fails, use the CPU fallback
            if (norm < 0) {
                printf("[%s] %s step=%d WARNING: Custom CUDA norm calculation failed, using CPU fallback\n", 
                      classname, TAG, step);
                norm = calculateNormCPUDouble(array, size);
                
                if (norm < 0) {
                    printf("[%s] %s step=%d ERROR: All norm calculation methods failed\n", 
                          classname, TAG, step);
                    return;
                }
            }
        }
        
        // Print the result
        printf("[%s] %s step=%d norm=%.12E\n", classname, TAG, step, norm);
        fflush(stdout);
    }
} 